import os
import cv2
import numpy as np
import pandas as pd
import glob
from pathlib import Path
import json
import sys
sys.path.append('..')
from utils.json import json_to_image, mask_to_json

# cut by training shape [h, w] = [640, 864]
std_h, std_w = 640, 864

def json_check(json_path, img_shape):
    with open(json_path, 'r') as f:
        content = f.read()
        data = json.loads(content)
    
    data['imageData'] = 'null'
    h, w = img_shape[:2]
    
    new_data_shapes= []
    for d in data['shapes']:
        if len(d['points']) < 3:
            continue
        d['label'] = d['label'].replace(' ', '')
        if (d['label']) == 6 and d['label'][-1] == '1':
            continue
        if (d['label']) == 6 and d['label'][-1] in ['2', '3']:
            d['label'] = d['label'][:-1]
        
        for i,p in enumerate(d['points']):  #[w, h]
            d['points'][i][0] = min(p[0], w)
            d['points'][i][1] = min(p[1], h)
        new_data_shapes.append(d)
    data['shapes'] = new_data_shapes

    with open(json_path, 'w') as f:
        b = json.dumps(data, indent=4)
        f.write(b)


def get_img_path(json_path, code, img_dst_root):
    img_name = str(Path(json_path).name)[:-4] + 'jpg'
    return os.path.join(img_dst_root, code, img_name)

def one_img_cut(img, mask):
    h, w = img.shape[:2]
    # mid_h = h//2
    # mid_w = w//2

    new_imgs, new_masks = [], []
    # for dh, dw in [[0, 0], [0, 1], [1, 1], [1, 0]]:
    #     new_imgs.append(img[dh*mid_h:(dh+1)*mid_h, dw*mid_w:(dw+1)*mid_w])
    #     new_masks.append(mask[dh*mid_h:(dh+1)*mid_h, dw*mid_w:(dw+1)*mid_w])
    
    x1, x2, x3, x4 = 0, 0, w - std_w, w - std_w  # for w
    y1, y2, y3, y4 = 0, h - std_h, 0, h - std_h  # for h
    for tx, ty in zip([x1, x2, x3, x4], [y1, y2, y3, y4]):
        new_imgs.append(img[ty:ty+std_h, tx:tx+std_w])
        new_masks.append(mask[ty:ty+std_h, tx:tx+std_w])

    
    return new_imgs, new_masks


# # json_to_image, mask_to_json
def img_cut(json_path, img_path, code, no_defect_code, dst_path):
    img_stem = str(Path(img_path).stem)
    
    img = cv2.imread(img_path, 1)
    json_check(json_path, img.shape)
    mask = json_to_image(img.shape, json_path)

    cutted_imgs, cutted_masks = one_img_cut(img, mask)
    for i, (c_img, c_mask) in enumerate(zip(cutted_imgs, cutted_masks)):
        # print('c mask shape: ', c_mask.shape)
        if 1 not in c_mask:
            tmp_code = 'X' + no_defect_code[1:]
        else:
            tmp_code = 'X' + code[1:] 

        img_name = img_stem + '_%d.jpg'%(i+1)
        mask_name = img_stem + '_%d.json'%(i+1)
        img_dst_path = os.path.join(dst_path, 'images', tmp_code)
        mask_dst_path = os.path.join(dst_path, 'annotations', tmp_code)

        os.makedirs(img_dst_path, exist_ok=True)
        cv2.imwrite(os.path.join(img_dst_path, img_name), c_img)

        if 1 in c_mask:
            os.makedirs(mask_dst_path, exist_ok=True)
            mask_json = mask_to_json(mask=c_mask, imagePath=img_name, label_name=tmp_code)
            with open(os.path.join(mask_dst_path, mask_name), 'w') as f:
                b = json.dumps(mask_json, indent=4)
                f.write(b)


# TDSBN TTSPT TCNFO  TTPRP TCSBP
def main():
    # parameters
    ################################################################ 
    img_path = r'/root/project/AutoRepair_T7/data/t7_training_data_new/all'
    json_path = r'/root/project/AutoRepair_T7/data/t7_training_data_new/all'
    dst_path = r'/root/project/AutoRepair_T7/data/t7_training_data_new/all_cutted_data_by_training_shape'

    # need_cut_codes = ['TDSBN', 'TTSPT', 'TCNFO', 'TTPRP', 'TCSBP']
    no_defect_code = 'TGXID'
    ################################################################

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    df = pd.DataFrame({'json':glob.glob(os.path.join(json_path, '*/*.json'))})
    df['code'] = df['json'].apply(lambda x:str(Path(x).parent.name))
    # df = df[df['code'].isin(need_cut_codes)].reset_index(drop=True)
    if is_pandarallel:
        df['img_path'] = df.parallel_apply(lambda x: get_img_path(x['json'], x['code'], img_path), axis=1)
    else:
        df['img_path'] = df.apply(lambda x: get_img_path(x['json'], x['code'], img_path), axis=1)
    
    print('len df: ', len(df))
    
    if is_pandarallel:
        df.parallel_apply(lambda x: img_cut(json_path=x['json'], img_path=x['img_path'], code=x['code'], no_defect_code=no_defect_code, dst_path=dst_path), axis=1)
    else:
        df.apply(lambda x: img_cut(json_path=x['json'], img_path=x['img_path'], code=x['code'], no_defect_code=no_defect_code, dst_path=dst_path), axis=1)


if __name__=='__main__':
    main()
